import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import ssl
from PIL import Image, ImageOps
import requests
from io import BytesIO

# 处理SSL证书问题
ssl._create_default_https_context = ssl._create_unverified_context

# 设备配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")

# 超参数（需要与训练时保持一致）
input_size = 784  # 28x28像素的图像
hidden_size = 128  # 隐藏层神经元数量
num_classes = 10  # 0-9共10个数字类别

# 定义与训练时相同的神经网络模型
class NeuralNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, num_classes)
        self.dropout = nn.Dropout(0.2)
    
    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.dropout(out)
        out = self.fc2(out)
        out = self.relu(out)
        out = self.dropout(out)
        out = self.fc3(out)
        return out

# 初始化模型
model = NeuralNet(input_size, hidden_size, num_classes).to(device)

# 加载训练好的模型权重
try:
    model.load_state_dict(torch.load('mnist_model.pth', map_location=device))
    print("成功加载模型")
except FileNotFoundError:
    print("未找到模型文件 'mnist_model.pth'，请先运行训练脚本")
    exit()

model.eval()
print("模型已准备好进行推理")

# 数据预处理
transform = transforms.Compose([transforms.ToTensor(), 
                                transforms.Normalize((0.1307,), (0.3081,))])

# 从测试集中获取一些样本进行推理演示
test_dataset = torchvision.datasets.MNIST(root='./data',
                                          train=False,
                                          transform=transform)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=1,
                                          shuffle=True)

# 显示图像并进行预测的函数
def imshow(img, title=None):
    img = img / 2 + 0.5  # 反向归一化
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)), cmap='gray')
    if title:
        plt.title(title)
    plt.show()

# 对单个图像进行预测
def predict_image(image_tensor):
    # 确保图像形状正确
    image_tensor = image_tensor.reshape(-1, 28*28).to(device)
    
    # 进行预测
    with torch.no_grad():
        output = model(image_tensor)
        probabilities = torch.nn.functional.softmax(output[0], dim=0)
        predicted_class = torch.argmax(output).item()
        confidence = probabilities[predicted_class].item()
    
    return predicted_class, confidence, probabilities

# 从测试集中选择一些样本进行推理
print("\n=== 手写数字识别推理演示 ===")
dataiter = iter(test_loader)

for i in range(5):  # 预测5个样本
    image, true_label = next(dataiter)
    
    # 进行预测
    predicted_class, confidence, probabilities = predict_image(image)
    
    # 显示结果
    print(f"\n样本 {i+1}:")
    print(f"真实标签: {true_label.item()}")
    print(f"预测标签: {predicted_class}")
    print(f"置信度: {confidence:.4f}")
    print("各类别概率:")
    for j in range(10):
        print(f"  {j}: {probabilities[j]:.4f}")
    
    # 显示图像
    title = f"真实: {true_label.item()}, 预测: {predicted_class} (置信度: {confidence:.2f})"
    imshow(torchvision.utils.make_grid(image), title)

# 创建一个函数来处理自定义图像
def predict_custom_image(image_path_or_url):
    """
    对自定义图像进行预测
    可以是本地文件路径或URL
    """
    try:
        # 加载图像
        if image_path_or_url.startswith('http'):
            # 从URL加载图像
            response = requests.get(image_path_or_url)
            img = Image.open(BytesIO(response.content)).convert('L')
        else:
            # 从本地文件加载图像
            img = Image.open(image_path_or_url).convert('L')
        
        # 预处理图像
        # 调整大小为28x28
        img = img.resize((28, 28))
        
        # 反转颜色（MNIST是白字黑底）
        img = ImageOps.invert(img)
        
        # 转换为tensor
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])
        img_tensor = transform(img)
        
        # 进行预测
        predicted_class, confidence, probabilities = predict_image(img_tensor.unsqueeze(0))
        
        # 显示结果
        print(f"\n=== 自定义图像预测结果 ===")
        print(f"预测标签: {predicted_class}")
        print(f"置信度: {confidence:.4f}")
        print("各类别概率:")
        for j in range(10):
            print(f"  {j}: {probabilities[j]:.4f}")
        
        # 显示图像
        plt.imshow(img, cmap='gray')
        plt.title(f'预测: {predicted_class} (置信度: {confidence:.2f})')
        plt.show()
        
    except Exception as e:
        print(f"处理自定义图像时出错: {e}")

print("\n推理演示完成")
print("\n要预测自定义图像，请调用 predict_custom_image('图像路径或URL') 函数")

# 使用示例（取消注释以使用）:
# predict_custom_image('path/to/your/image.png')
# predict_custom_image('https://example.com/your_digit_image.png')